# -*- coding: utf-8 -*-
#!/bin/python3
# 根据模板生成Tansible的action文件，用于批量进行防火墙配置检验
# python3 gen_firewall.py 运行
import os
import yaml

# 测试，num的值表示测试读取配置行数，0表示输出所有
num = 0

mac_file = 'firewall.txt'
tpl_file = 'tmpl.txt'
out_file = 'out.txt'
 
HOSTS_CFG = 'hosts.yaml'

# 定义列标题映射，firewall.txt第一行为标题行，一定要按中文名给定，顺序可以随意
COLUMNS = {
    "源主机名": "hostname",
    "目的IP地址": "ip",
    "目的端口": "port",
    "备注": "remark",
}

def parse_header(header_line):
    """解析列标题，返回列名到索引的映射"""
    headers = header_line.strip().split()
    column_map = {}
    for idx, header in enumerate(headers):
        if header in COLUMNS:
            column_map[COLUMNS[header]] = idx
    return column_map

def parse_line(line, column_map):
    """解析数据行，返回字典形式的数据"""
    parts = line.strip().split()
    data = {}
    for key, idx in column_map.items():
        if idx < len(parts):
            data[key] = parts[idx].strip()
        else:
            data[key] = ""
    return data

def get_fw_list(mac_file):
    fw_list = []

    with open(mac_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()
        global num
        if num == 0:
            num = len(lines)
        
        # 解析列标题
        column_map = parse_header(lines[0])
        
        for line in lines[1:num + 1]:  # 从第二行开始读取数据
            data = parse_line(line, column_map)
            hostname = data.get("hostname", "")
            ip = data.get("ip", "")
            port = data.get("port", "")
            remark = data.get("remark", "")
            
            # hostname处理
            hostname_lst = []
            if "/" in hostname:
                hostname_ss = hostname.split("/")
                hostname_len = len(hostname_ss[1])
                hostname_name = hostname_ss[0][:-hostname_len]
                for m,hostname_num in enumerate(hostname_ss):
                    if m == 0:
                        hostname_lst.append(hostname_ss[m])
                    else:
                        hostname_lst.append(f"{hostname_name}{hostname_num}")
            else:
                hostname_lst.append(hostname)

            # ip处理
            ip_lst = []
            #ip格式：s1.s2.s3.s4_1/s4_2
            if "/" in ip:
                ip_all = ip.split(".")
                ip_s4s = ip_all[3].split("/")
                for ip_s4 in ip_s4s:
                    ip_item = f"{ip_all[0]}.{ip_all[1]}.{ip_all[2]}.{ip_s4}"
                    ip_lst.append(ip_item)
            else:
                ip_lst.append(ip)
           
            # 全部数据push到fw_list列表中
            for hostname_item in hostname_lst:
                for ip_item in ip_lst: 
                    fw_item = {"remark":remark, "hostname":hostname_item, "ip":ip_item, "port":port}
                    fw_list.append(fw_item)

    return fw_list

def render_tpl(tpl_content, hosts_data, fw_item):
    remark = fw_item["remark"]   
    hostname = fw_item["hostname"]  
    ip = fw_item["ip"] 
    port = fw_item["port"] 

    # 如果HOSTS_CFG 存在则查找第一个替换hostname,    
    if hosts_data is not None:                             
        
        # 遍历HOST下的所有key，查找第一个包含指定字符的key
        host_dict = hosts_data['HOST']
        for key in host_dict.keys():
            # 确保key是字符串类型（避免非字符串key报错）
            if isinstance(key, str) and hostname in key:
                content = tpl_content.replace('{REMARK}', remark).replace('{HOSTNAME}',key).replace('{IP}', ip).replace('{PORT}', port)
                return content
        
        # 未找到匹配的key
        print(f"未在HOST层级下找到包含 '{hostname}' 的key")
        return None
       
    else:
        content = tpl_content.replace('{REMARK}', remark).replace('{HOSTNAME}',hostname+"*").replace('{IP}', ip).replace('{PORT}', port)

    return content

def gen_result(out_file, tpl_file, fw_list):
    if os.path.exists(out_file):
        os.remove(out_file)
   
   # 读取模板文件
    tpl_content = ''
    with open(tpl_file, 'r', encoding='utf-8') as f:
         tpl_content = f.read()
    
    hosts_data = None
    if os.path.exists(HOSTS_CFG):
        # 读取并解析YAML文件
        with open(HOSTS_CFG, 'r', encoding='utf-8') as y:
            hosts_data = yaml.safe_load(y)  

    for idx, fw_item in enumerate(fw_list):  
        hostname = fw_item["hostname"]  
        ip = fw_item["ip"] 
        port = fw_item["port"] 

        content = render_tpl(tpl_content, hosts_data, fw_item)
            
        with open(out_file, "a", encoding='utf-8') as f:
            print(f'生成配置文件 {idx}: 主机 {hostname:20s} IP {ip:15s} PORT {port:15s}')
            f.write(content)


fw_list = get_fw_list(mac_file)
gen_result(out_file, tpl_file, fw_list)